import tkinter as tk
from tkinter import ttk, filedialog, scrolledtext, filedialog
import threading
import queue
import os
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException
from selenium.webdriver.common.keys import Keys
import time
import sys
import logging
import re
from datetime import datetime
from functools import wraps
import requests
from bs4 import BeautifulSoup

class RetryableError(Exception):
    """可重试的错误"""
    pass


def retry_on_failure(max_retries=3, delay=2):
    """重试装饰器"""

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            for attempt in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except RetryableError as e:
                    last_exception = e
                    if attempt < max_retries - 1:
                        logging.warning(f"第{attempt + 1}次尝试失败: {str(e)}, {delay}秒后重试...")
                        time.sleep(delay)
            logging.error(f"在{max_retries}次尝试后失败: {str(last_exception)}")
            raise last_exception

        return wrapper

    return decorator


class RedirectText:
    def __init__(self, text_widget, logger):
        self.text_widget = text_widget
        self.queue = queue.Queue()
        self.update_timer = None
        self.logger = logger

    def write(self, string):
        self.queue.put(string)
        if self.update_timer is None:
            self.update_timer = self.text_widget.after(100, self.update_text)
        # 同时写入日志
        string = string.strip()
        if string:
            self.logger.info(string)

    def update_text(self):
        while not self.queue.empty():
            string = self.queue.get()
            self.text_widget.configure(state='normal')
            self.text_widget.insert(tk.END, string)
            self.text_widget.see(tk.END)
            self.text_widget.configure(state='disabled')
        self.update_timer = None

    def flush(self):
        pass


def clean_suggestion(suggestion):
    """清理建议文本"""
    # 移除特殊字符但保留中文、英文、数字和基本标点
    cleaned = re.sub(r'[^\w\s\u4e00-\u9fff.,!?，。！？]', '', suggestion)
    # 统一空格
    cleaned = ' '.join(cleaned.split())
    # 移除前后的空格和标点
    cleaned = cleaned.strip('.,!?，。！？ ')
    return cleaned


class RedirectText:
    def __init__(self, text_widget, logger):
        self.text_widget = text_widget
        self.logger = logger

    def write(self, string):
        self.text_widget.configure(state='normal')
        self.text_widget.insert(tk.END, string)
        self.text_widget.see(tk.END)
        self.text_widget.configure(state='disabled')
        self.logger.info(string.strip())

    def flush(self):
        pass


def parse_input_file(file_path):
    """解析带有类别的输入文件
    返回: [{'category': '类别', 'google_us': 'Google US关键词', 'google_cn': 'Google CN关键词'}]
    """
    keywords = []
    current_category = None

    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:  # 跳过空行
                continue
            if line.startswith('# '):  # 类别标题
                current_category = line[2:]
            elif line.startswith('## '):  # 搜索关键词对
                # 分割Google US和Google CN关键词
                kw_pair = line[3:].split(' | ')
                if len(kw_pair) == 2:
                    keywords.append({
                        'category': current_category,
                        'google_us': kw_pair[0].strip(),
                        'google_cn': kw_pair[1].strip()
                    })
    return keywords


class KeywordScraperGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("Google Search Suggestion Scraper")
        self.root.geometry("800x700")

        # 设置日志
        self.setup_logging()

        self.main_frame = ttk.Frame(root, padding="10")
        self.main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))

        self.create_platform_selection()
        self.create_file_selection()
        self.create_log_area()
        self.create_progress_area()
        self.create_control_buttons()

        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        self.main_frame.columnconfigure(1, weight=1)

        self.is_running = False
        self.current_progress = 0
        self.total_progress = 0

    def log(self, message, level="INFO", detail_level="basic"):
        """
        统一的日志记录函数
        detail_level: 日志详细程度 (basic/detailed)
        """
        # 如果是详细日志，但当前设置是基础级别，则不输出
        if detail_level == "detailed" and self.log_level_var.get() == "basic":
            return

        if level == "ERROR":
            prefix = "× "
        elif level == "WARNING":
            prefix = "! "
        elif level == "SUCCESS":
            prefix = "✓ "
        else:
            prefix = ""

        # 确保消息后有换行符
        if not message.endswith('\n'):
            message += '\n'

        print(prefix + message, end='')

    def setup_logging(self):
        """设置日志系统"""
        log_dir = "logs"
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)

        log_file = os.path.join(log_dir, f"scraper_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file, encoding='utf-8'),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger(__name__)

    def create_platform_selection(self):
        platform_frame = ttk.LabelFrame(self.main_frame, text="设置", padding="5")
        platform_frame.grid(row=0, column=0, columnspan=3, sticky=(tk.W, tk.E), pady=5)

        # 平台选择区域
        platform_select_frame = ttk.Frame(platform_frame)
        platform_select_frame.pack(fill=tk.X, pady=5)

        self.google_us_var = tk.BooleanVar(value=True)
        self.google_cn_var = tk.BooleanVar(value=True)
        ttk.Checkbutton(platform_select_frame, text="Google US", variable=self.google_us_var).pack(side=tk.LEFT,
                                                                                                   padx=10)
        ttk.Checkbutton(platform_select_frame, text="Google CN", variable=self.google_cn_var).pack(side=tk.LEFT,
                                                                                                   padx=10)

        # 输出格式选择
        self.output_format_var = tk.StringVar(value="single")
        ttk.Radiobutton(platform_select_frame, text="合并输出", variable=self.output_format_var,
                        value="single").pack(side=tk.LEFT, padx=10)
        ttk.Radiobutton(platform_select_frame, text="分别输出", variable=self.output_format_var,
                        value="separate").pack(side=tk.LEFT, padx=10)

        # 浏览器显示选项
        self.show_browser_var = tk.BooleanVar(value=False)
        browser_frame = ttk.Frame(platform_frame)
        browser_frame.pack(fill=tk.X, pady=5)
        ttk.Checkbutton(browser_frame, text="显示浏览器（如获取失败请开启此选项）",
                        variable=self.show_browser_var).pack(side=tk.LEFT, padx=10)

        # 日志级别选择区域
        log_level_frame = ttk.Frame(platform_frame)
        log_level_frame.pack(fill=tk.X, pady=5)

        ttk.Label(log_level_frame, text="日志级别:").pack(side=tk.LEFT, padx=5)
        self.log_level_var = tk.StringVar(value="detailed")
        ttk.Radiobutton(log_level_frame, text="基础", variable=self.log_level_var,
                        value="basic").pack(side=tk.LEFT, padx=10)
        ttk.Radiobutton(log_level_frame, text="详细", variable=self.log_level_var,
                        value="detailed").pack(side=tk.LEFT, padx=10)

    def create_file_selection(self):
        # Input file selection
        ttk.Label(self.main_frame, text="输入文件:").grid(row=1, column=0, sticky=tk.W, pady=5)
        self.input_path = tk.StringVar()
        self.input_entry = ttk.Entry(self.main_frame, textvariable=self.input_path)
        self.input_entry.grid(row=1, column=1, sticky=(tk.W, tk.E), padx=5)
        ttk.Button(self.main_frame, text="浏览", command=self.browse_input).grid(row=1, column=2)

        # Output file frame
        self.output_frame = ttk.Frame(self.main_frame)
        self.output_frame.grid(row=2, column=0, columnspan=3, sticky=(tk.W, tk.E), pady=5)

        # Single output
        self.single_output_frame = ttk.Frame(self.output_frame)
        ttk.Label(self.single_output_frame, text="输出文件:").pack(side=tk.LEFT)
        self.output_path = tk.StringVar()
        ttk.Entry(self.single_output_frame, textvariable=self.output_path).pack(side=tk.LEFT, fill=tk.X, expand=True,
                                                                                padx=5)
        ttk.Button(self.single_output_frame, text="浏览", command=self.browse_output).pack(side=tk.LEFT)

        # Separate outputs
        self.separate_output_frame = ttk.Frame(self.output_frame)
        # Google US output
        ttk.Label(self.separate_output_frame, text="Google US输出:").pack(side=tk.LEFT)
        self.google_us_output_path = tk.StringVar()
        ttk.Entry(self.separate_output_frame, textvariable=self.google_us_output_path).pack(side=tk.LEFT, fill=tk.X,
                                                                                            expand=True, padx=5)
        ttk.Button(self.separate_output_frame, text="浏览",
                   command=lambda: self.browse_output_separate('google_us')).pack(side=tk.LEFT)
        # Google CN output
        ttk.Label(self.separate_output_frame, text="Google CN输出:").pack(side=tk.LEFT, padx=(10, 0))
        self.google_cn_output_path = tk.StringVar()
        ttk.Entry(self.separate_output_frame, textvariable=self.google_cn_output_path).pack(side=tk.LEFT, fill=tk.X,
                                                                                            expand=True, padx=5)
        ttk.Button(self.separate_output_frame, text="浏览",
                   command=lambda: self.browse_output_separate('google_cn')).pack(side=tk.LEFT)

        # Update output frame based on selection
        self.output_format_var.trace('w', self.update_output_frame)
        self.update_output_frame()

    def create_log_area(self):
        """创建日志显示区域"""
        log_frame = ttk.LabelFrame(self.main_frame, text="运行日志", padding="5")
        log_frame.grid(row=3, column=0, columnspan=3, sticky=(tk.W, tk.E, tk.N, tk.S), pady=10)

        # 创建和配置日志文本框
        self.log_text = scrolledtext.ScrolledText(log_frame, height=20, width=80, wrap=tk.WORD)
        self.log_text.pack(fill=tk.BOTH, expand=True)
        self.log_text.configure(state='disabled')

        # 配置标准输出重定向
        sys.stdout = RedirectText(self.log_text, self.logger)

    def create_progress_area(self):
        """创建进度显示区域"""
        progress_frame = ttk.Frame(self.main_frame)
        progress_frame.grid(row=4, column=0, columnspan=3, sticky=(tk.W, tk.E), pady=5)

        self.progress_var = tk.StringVar(value="准备就绪")
        self.progress_label = ttk.Label(progress_frame, textvariable=self.progress_var)
        self.progress_label.pack(pady=2)

        self.progress_bar = ttk.Progressbar(progress_frame, mode='determinate')
        self.progress_bar.pack(fill=tk.X, pady=2)

    def create_control_buttons(self):
        """创建控制按钮"""
        button_frame = ttk.Frame(self.main_frame)
        button_frame.grid(row=5, column=0, columnspan=3, pady=10)

        self.start_button = ttk.Button(button_frame, text="开始运行", command=self.start_scraping)
        self.start_button.pack(side=tk.LEFT, padx=5)

        self.clear_button = ttk.Button(button_frame, text="清除日志", command=self.clear_log)
        self.clear_button.pack(side=tk.LEFT, padx=5)

    def update_output_frame(self, *args):
        """更新输出文件选择框的显示"""
        if self.output_format_var.get() == "single":
            self.separate_output_frame.grid_remove()
            self.single_output_frame.grid(row=0, column=0, sticky=(tk.W, tk.E))
        else:
            self.single_output_frame.grid_remove()
            self.separate_output_frame.grid(row=0, column=0, sticky=(tk.W, tk.E))

    def browse_input(self):
        """选择输入文件"""
        filename = filedialog.askopenfilename(
            title="选择输入文件",
            filetypes=[("Text Files", "*.txt"), ("All Files", "*.*")]
        )
        if filename:
            self.input_path.set(filename)
            self.log(f"选择输入文件: {filename}", detail_level="detailed")

    def browse_output(self):
        """选择输出文件"""
        filename = filedialog.asksaveasfilename(
            title="选择输出文件",
            defaultextension=".txt",
            filetypes=[("Text Files", "*.txt"), ("All Files", "*.*")]
        )
        if filename:
            self.output_path.set(filename)
            self.log(f"选择输出文件: {filename}", detail_level="detailed")

    def browse_output_separate(self, platform):
        """选择分别输出的文件"""
        filename = filedialog.asksaveasfilename(
            title=f"选择{platform}输出文件",
            defaultextension=".txt",
            filetypes=[("Text Files", "*.txt"), ("All Files", "*.*")]
        )
        if filename:
            if platform == 'google_us':
                self.google_us_output_path.set(filename)
            else:
                self.google_cn_output_path.set(filename)
            self.log(f"选择{platform}输出文件: {filename}", detail_level="basic")

    def clear_log(self):
        """清除日志显示"""
        self.log_text.configure(state='normal')
        self.log_text.delete(1.0, tk.END)
        self.log_text.configure(state='disabled')
        self.log("清除日志显示", detail_level="detailed")

    def validate_inputs(self):
        """验证输入是否有效"""
        if not self.input_path.get():
            self.log("× 请选择输入文件！", detail_level="basic")
            return False

        if not os.path.exists(self.input_path.get()):
            self.log("× 输入文件不存在！", detail_level="basic")
            return False

        if not self.google_us_var.get() and not self.google_cn_var.get():
            self.log("× 请至少选择一个平台！", detail_level="basic")
            return False

        if self.output_format_var.get() == "single":
            if not self.output_path.get():
                self.log("× 请选择输出文件！", detail_level="basic")
                return False
        else:
            if self.google_us_var.get() and not self.google_us_output_path.get():
                self.log("× 请选择Google US输出文件！", detail_level="basic")
                return False
            if self.google_cn_var.get() and not self.google_cn_output_path.get():
                self.log("× 请选择Google CN输出文件！", detail_level="basic")
                return False

        return True

    def update_progress(self, current, total, message=None):
        """更新进度显示"""
        self.current_progress = current
        self.total_progress = total
        progress = (current / total * 100) if total > 0 else 0

        self.progress_bar['value'] = progress
        if message:
            self.progress_var.set(f"{message} - {current}/{total} ({progress:.1f}%)")
        else:
            self.progress_var.set(f"进度: {current}/{total} ({progress:.1f}%)")

    def start_scraping(self):
        """开始爬取关键词"""
        if self.is_running:
            self.log("任务已在运行中...", detail_level="basic")
            return

        if not self.validate_inputs():
            return

        self.is_running = True
        self.start_button.configure(state='disabled')

        # 重置进度条
        self.progress_bar['value'] = 0
        self.progress_var.set("准备开始...")

        # 在新线程中运行爬虫
        thread = threading.Thread(target=self.run_scraper)
        thread.daemon = True
        thread.start()

    def run_scraper(self):
        """运行爬虫主流程"""
        try:
            self.log(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] 开始运行...", detail_level="basic")

            results = {}
            if self.google_us_var.get():
                self.log("\n开始获取Google US关键词建议...", detail_level="basic")
                results['google_us'] = self.get_google_suggestions('us')

            if self.google_cn_var.get():
                self.log("\n开始获取Google CN关键词建议...", detail_level="basic")
                results['google_cn'] = self.get_google_suggestions('cn')

            self.save_results(results)

            self.log(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] 运行完成!", detail_level="basic")
            self.progress_var.set("任务完成")
        except Exception as e:
            self.log(f"× 运行出错: {str(e)}", detail_level="basic")
            self.progress_var.set("运行出错")
        finally:
            self.is_running = False
            self.root.after(0, lambda: self.start_button.configure(state='normal'))

    def get_browser_options(self):
        """获取浏览器选项"""
        options = webdriver.ChromeOptions()

        # 如果用户未选择显示浏览器时使用无头模式
        if not self.show_browser_var.get():
            options.add_argument('--headless=new')

        options.add_argument('--window-size=1920,1080')
        options.add_argument('--disable-gpu')
        options.add_argument('--no-sandbox')
        options.add_argument('--disable-dev-shm-usage')
        options.add_experimental_option('excludeSwitches', ['enable-logging'])

        return options

    def parse_google_suggestions(self, html_content):
        """解析Google搜索建议HTML"""
        soup = BeautifulSoup(html_content, 'html.parser')
        suggestions = []

        # 查找所有搜索建议元素
        suggestion_elements = soup.select('li.sbct')

        for element in suggestion_elements:
            # 获取建议文本
            suggestion_text_element = element.select_one('div.wM6W7d')
            if suggestion_text_element:
                # 获取纯文本内容并去除HTML标签
                suggestion_text = suggestion_text_element.get_text(strip=True)
                if suggestion_text and suggestion_text not in suggestions:
                    suggestions.append(suggestion_text)

        return suggestions

    def get_suggestions_for_keyword(self, driver, keyword, region='us'):
        """获取单个关键词的建议"""
        try:
            # 设置Google搜索URL (根据地区选择不同的Google域名)
            if region == 'us':
                url = 'https://www.google.com/'
            else:  # 中国
                url = 'https://www.google.com.hk/'

            driver.get(url)
            time.sleep(2)

            # 查找搜索框并输入关键词
            search_box = WebDriverWait(driver, 10).until(
                EC.presence_of_element_located((By.NAME, "q"))
            )

            search_box.clear()
            search_box.send_keys(keyword)
            time.sleep(1.5)  # 等待搜索建议加载

            # 获取页面HTML
            html_content = driver.page_source

            # 解析搜索建议
            suggestions = self.parse_google_suggestions(html_content)

            if suggestions:
                self.log(f"✓ 关键词 '{keyword}' 找到 {len(suggestions)} 个建议", detail_level="detailed")
                for suggestion in suggestions:
                    self.log(f"  - {suggestion}", detail_level="detailed")
            else:
                self.log(f"! 关键词 '{keyword}' 未找到建议", detail_level="detailed")

            return suggestions

        except Exception as e:
            self.log(f"× 获取 '{keyword}' 的建议时出错: {str(e)}", detail_level="basic")
            return []

    def get_google_suggestions(self, region='us'):
        """获取Google关键词建议"""
        suggestions = {}
        region_name = "Google US" if region == 'us' else "Google CN"
        self.log(f"\n===== 开始获取{region_name}关键词建议 =====", detail_level="basic")

        options = self.get_browser_options()
        driver = webdriver.Chrome(options=options)

        try:
            # 读取关键词
            keywords = parse_input_file(self.input_path.get())
            total = len(keywords)
            self.update_progress(0, total, f"处理{region_name}关键词")

            keyword_field = 'google_us' if region == 'us' else 'google_cn'

            # 处理每个关键词
            for index, keyword in enumerate(keywords, 1):
                try:
                    kw = keyword[keyword_field]
                    self.log(f"\n----- 处理{region_name}关键词 ({index}/{total}): {kw} -----", detail_level="basic")
                    suggestions[kw] = self.get_suggestions_for_keyword(driver, kw, region)
                    self.update_progress(index, total, f"处理{region_name}关键词")
                    time.sleep(1)  # 避免频繁请求
                except Exception as e:
                    self.log(f"处理关键词 '{kw}' 时出错: {str(e)}", detail_level="basic")
                    continue

        finally:
            self.log(f"\n关闭Chrome浏览器({region_name})...", detail_level="detailed")
            driver.quit()
            self.log(f"===== {region_name}关键词获取完成 =====\n", detail_level="basic")

        return suggestions

    def save_results(self, results):
        """保存结果到文件"""
        if self.output_format_var.get() == "single":
            with open(self.output_path.get(), 'w', encoding='utf-8') as f:
                current_category = None
                keywords = parse_input_file(self.input_path.get())

                for keyword in keywords:
                    # 如果类别变化，写入新类别
                    if keyword['category'] != current_category:
                        current_category = keyword['category']
                        f.write(f"# {current_category}\n")

                    # 写入关键词
                    f.write(f"## {keyword['google_cn']}\n")

                    # 获取两个平台的建议
                    us_suggestions = results.get('google_us', {}).get(keyword['google_us'], [])
                    cn_suggestions = results.get('google_cn', {}).get(keyword['google_cn'], [])

                    # 取建议数量的最小值，确保一一对应
                    min_suggestions = min(len(us_suggestions), len(cn_suggestions))

                    # 写入配对的建议
                    for i in range(min_suggestions):
                        f.write(f"{us_suggestions[i]} | {cn_suggestions[i]}\n")
                    f.write("\n")
        else:
            # 分别输出到不同文件
            if 'google_us' in results and self.google_us_var.get():
                with open(self.google_us_output_path.get(), 'w', encoding='utf-8') as f:
                    current_category = None
                    keywords = parse_input_file(self.input_path.get())

                    for keyword in keywords:
                        if keyword['category'] != current_category:
                            current_category = keyword['category']
                            f.write(f"# {current_category}\n")
                        # 写入Google US关键词
                        f.write(f"## {keyword['google_us']}\n")
                        if keyword['google_us'] in results['google_us']:
                            for suggestion in results['google_us'][keyword['google_us']]:
                                f.write(f"{suggestion}\n")
                        f.write("\n")

            if 'google_cn' in results and self.google_cn_var.get():
                with open(self.google_cn_output_path.get(), 'w', encoding='utf-8') as f:
                    current_category = None
                    keywords = parse_input_file(self.input_path.get())

                    for keyword in keywords:
                        if keyword['category'] != current_category:
                            current_category = keyword['category']
                            f.write(f"# {current_category}\n")
                        # 写入Google CN关键词
                        f.write(f"## {keyword['google_cn']}\n")
                        if keyword['google_cn'] in results['google_cn']:
                            for suggestion in results['google_cn'][keyword['google_cn']]:
                                f.write(f"{suggestion}\n")
                        f.write("\n")


def main():
    root = tk.Tk()
    app = KeywordScraperGUI(root)
    root.mainloop()


if __name__ == "__main__":
    main()


